package com.wxc.impl;

import com.wxc.MyMap;
import com.wxc.rbtree.IComparator;
import org.jetbrains.annotations.NotNull;

import java.util.LinkedList;

/**
 * @author wxc
 * @date 2021年08月16日 9:51 PM
 * key 不允许空 value允许空
 * key 需要具备可比较性 value不需要
 */
public class TreeMap<K, V> implements MyMap<K, V> {

	private static final boolean RED = false;
	private static final boolean BLACK = true;
	private int size;
	private Entry<K, V> root;
	private IComparator<K> comparator;

	public TreeMap(IComparator<K> comparator) {
		this.comparator = comparator;
	}

	public TreeMap() {
	}

	@Override
	public int size() {
		return size;
	}

	@Override
	public boolean isEmpty() {
		return size() == 0;
	}

	@Override
	public void clear() {
		size = 0;
		root = null;
	}

	@Override
	public V put(K key, V value) {
		elementNotNullCheck(key);
		if (root == null) {
			root = createEntry(key, value, null, null, null);
			size++;
			afterPut(root);
			return null;
		}

		Entry<K, V> node = root;
		Entry<K, V> parent = root;
		int cmp = 0;
		while (node != null) {
			cmp = compare(key, node.key);
			// 父结点
			parent = node;
			if (cmp > 0) {
				node = node.right;
			} else if (cmp < 0) {
				node = node.left;
			} else {
				// 相等
				final V OLD_VALUE = node.value;
				node.value = value;
				node.key = key;
				return OLD_VALUE;
			}
		}
		Entry<K, V> newNode = createEntry(key, value, null, null, parent);
		if (cmp > 0) parent.right = newNode;
		else parent.left = newNode;
		size++;
		// 旋转
		afterPut(newNode);
		return null;
	}

	@Override
	public V get(K key) {
		Entry<K, V> entry = getEntry(key);
		return entry == null ? null : entry.value;
	}

	@Override
	public V remove(K key) {
		return remove(getEntry(key));
	}

	@Override
	public boolean containsKey(K key) {
		return getEntry(key) != null;
	}

	@Override
	public boolean containsValue(V value) {
		if (root == null) return false;
		LinkedList<Entry<K, V>> queue = new LinkedList<>();
		queue.offer(root);
		while(!queue.isEmpty()) {
			Entry<K, V> node = queue.poll();
//			if(value == null && node.value == null) return true;
//			else if(node.value != null && node.value.equals(value)) return true;
			// 可以刚开始就判断value是不是为空，这样子的话每次循环都要进到这个方法里判断两次
			if(valEquals(value,node.value)) return true;

			if(node.left != null) queue.offer(node.left);
			if(node.right != null) queue.offer(node.right);
		}
		return false;
	}

	@Override
	public void traversal(Visitor<K, V> visitor) {
		if(visitor == null) return;
		traversal(root,visitor);
	}

	private void traversal(Entry<K, V> node,Visitor<K, V> visitor){
		if (node == null || visitor == null || visitor.stop) return;
		traversal(node.left, visitor);

		if (visitor.stop) return;
		visitor.stop = visitor.visit(node.key,node.value);

		traversal(node.right, visitor);
	}

	private boolean valEquals(V v1,V v2) {
		// 好绝
		return v1 == null ? v2 == null : v2.equals(v1);
	}

	private Entry<K, V> getEntry(K key) {
		Entry<K, V> node = root;
		int cmp = 0;
		while (node != null) {
			cmp = compare(key, node.key);
			if (cmp > 0) {
				node = node.right;
			} else if (cmp < 0) {
				node = node.left;
			} else {
				return node;
			}
		}
		return null;
	}

	private V remove(Entry<K, V> node) {
		if (node == null) return null;
		size--;
		// 度为2
		if (node.hasFullChild()) {
			// 后继
			Entry<K, V> nextNode = getSuccessor(node);
			// 覆盖
			node.key = nextNode.key;
			node.value = nextNode.value;
			// 删除 nextNode 这步太妙了
			node = nextNode;
		}
		// 度为0
		if (node.isLeaf()) {
			// 根结点
			if (node.parent == null) {
				root = null;
			}
			// 不是根结点
			else {
				if (node.parent.left == node) node.parent.left = null;
				else node.parent.right = null;
			}
			V value = node.value;
			// 删除的是叶子结点 所以取代被删除的结点是null
			afterRemove(node, null);
			return value;
		}
		// 度为1
		else {
			// 左边取代node
			if (node.left != null) {
				node.left.parent = node.parent;
				//          A           A
				//       (B)    ->    C
				//      C
				if (node.parent == null) {
					root = node.left;
					V value = node.value;
					afterRemove(node, root);
					return value;
				}
				if (node == node.parent.left) node.parent.left = node.left;
					//          A              A
					//              (B)    ->    C
					//            C
				else node.parent.right = node.left;

				V value = node.value;
				afterRemove(node, node.left);
				return value;
			}
			// 右边取代node
			else {
				node.right.parent = node.parent;

				if (node.parent == null) {
					root = node.right;
					V value = node.value;
					afterRemove(node, root);
					return value;
				} else if (node == node.parent.left) node.parent.left = node.right;
				else node.parent.right = node.right;
				V value = node.value;
				afterRemove(node, node.right);
				return value;
			}
		}
	}

	private void afterRemove(Entry<K, V> node, Entry<K, V> replaceNode) {
		// 1/2/3 Red
		if (isRed(node)) return;
			// 被删除的结点是Black
		else {
			// 用于替代的子节点是红色 null也是黑色
			// 由于node是黑色 所以要染成黑色
			if (isRed(replaceNode)) {
				dyeingBlack(replaceNode);
				return;
			}
			// 删除的是黑色叶子结点
			else {
				// 如果删除的是根节点
				Entry<K, V> parentNode = node.parent;
				if (parentNode == null) return;
					// 不是根节点
				else {
					/**
					 * 			 80
					 * 			/  \
					 * 		   /    \
					 * 		  76    88
					 * 		   \
					 * 		   78(红)
					 */

					/**
					 * 		   	 80
					 * 			/  \
					 * 		   /    \
					 * 		  76    88
					 * 		 /
					 * 		72(红)
					 */

					/**
					 * 		     80
					 * 			/  \
					 * 		   /    \
					 * 		  76    88
					 * 		 /  \
					 * 	(红)72  78(红)
					 */
					// 拿到兄弟结点
					// 这样写是有问题的
					// 因为 ndoe是叶子结点 在isLeftChildNode()这个函数中parent.left == this;永远为false
					// 因为在删除叶子结点remove()方法逻辑中有node.parent.left = null;
					// Node<E> siblingNode = node.getSibling();
					// 找到被删除的是左边还是右边
					// node.isLeftChildNode() 是为了解决父结点下溢的情况，parent先没断掉
					boolean left = parentNode.left == null || node.isLeftChildNode();
					Entry<K, V> siblingNode = left ? parentNode.right : parentNode.left;
					// 被删除的结点在左边
					if (left) {
						// 兄弟结点是红色
						if (isRed(siblingNode)) {
							// 兄弟染成黑色
							dyeingBlack(siblingNode);
							// 父结点染成红色
							dyeingRed(parentNode);
							// 对父结点进行右旋
							rotateLeft(parentNode);
							// 旋转完父结点变成红色 现在兄弟结点变成黑色
							// 和删除一个度为2的子节点的代码形式很像
							// 删除 nextNode 这步太妙了
							// node = nextNode;
							// 更换兄弟
							siblingNode = parentNode.right;
						}
						// 兄弟结点是黑色
						// 看我的黑兄弟能不能借 -> 能借的话兄弟要有红子节点
						if (isBlack(siblingNode.left) && isBlack(siblingNode.right)) {
							// 父结点要向下合并
							boolean parentNodeBlack = isBlack(parentNode);
							dyeingBlack(parentNode);
							dyeingRed(siblingNode);
							if (parentNodeBlack) {
								afterRemove(parentNode, null);
							}
						}
						// 至少有一红 向兄弟结点借
						else {
							if (isBlack(siblingNode.right)) {
								rotateRight(siblingNode);
								// 改变sibling
								siblingNode = parentNode.right;
							}
							// 先染色 再旋转
							dyeing(siblingNode, colorOf(parentNode));
							dyeingBlack(siblingNode.right);
							dyeingBlack(parentNode);

							// 统一右旋
							rotateLeft(parentNode);

						}
					}
					// 被删除结点在右边
					else {
						// 兄弟结点是红色
						if (isRed(siblingNode)) {
							// 兄弟染成黑色
							dyeingBlack(siblingNode);
							// 父结点染成红色
							dyeingRed(parentNode);
							// 对父结点进行右旋
							rotateRight(parentNode);
							// 旋转完父结点变成红色 现在兄弟结点变成黑色
							// 和删除一个度为2的子节点的代码形式很像
							// 删除 nextNode 这步太妙了
							// node = nextNode;
							// 更换兄弟
							siblingNode = parentNode.left;
						}
						// 兄弟结点是黑色
						// 看我的黑兄弟能不能借 -> 能借的话兄弟要有红子节点
						if (isBlack(siblingNode.left) && isBlack(siblingNode.right)) {
							// 父结点要向下合并
							boolean parentNodeBlack = isBlack(parentNode);
							dyeingBlack(parentNode);
							dyeingRed(siblingNode);
							if (parentNodeBlack) {
								afterRemove(parentNode, null);
							}
						}
						// 至少有一红 向兄弟结点借
						else {
							if (isBlack(siblingNode.left)) {
								rotateLeft(siblingNode);
								// 改变sibling
								siblingNode = parentNode.left;
							}
							// 先染色 再旋转
							dyeing(siblingNode, colorOf(parentNode));
							dyeingBlack(siblingNode.left);
							dyeingBlack(parentNode);

							// 统一右旋
							rotateRight(parentNode);

						}
					}
				}
			}
		}
	}

	private Entry<K, V> getSuccessor(Entry<K, V> node) {
		if (node == null) return null;
		// 说明是根结点
		if (node.right == null && node.parent == null) return null;
		// 一直往上走，如果父结点的是父父结点的右子结点,那么pre就是父父结点
		// 具体可以看视频
		Entry<K, V> pre = null;
		if (node.right == null && node.parent != null) {
			pre = node;
			while (pre != null && pre.parent != null) {
				if (pre.parent.left == pre) return pre.parent;
				pre = node.parent;
			}
		}
		// 先往右 再一直往左
		else if (node.right != null) {
			pre = node.right;
			while (pre.left != null) pre = pre.left;
			return pre;
		}
		return pre;
	}

	private int compare(K key1, K key2) {
		if (comparator != null) return comparator.compareTo(key1, key2);
		// 如果使用自定义比较器的话，那泛型E是没必要实现Comparable接口的，但是在这里就需要强制转型，因为上面没走进去
		if (key1 instanceof Comparable) return ((Comparable<K>) key1).compareTo(key2);
		throw new ClassCastException("not implements Comparable interface");
	}

	private void afterPut(Entry<K, V> node) {
		Entry<K, V> parentNode = node.parent;
		// 没有父结点
		if (parentNode == null) {
			dyeingBlack(node);
			return;
		}
		// 父结点是黑色
		if (isBlack(parentNode)) return;
			// 父结点是红色
		else {
			// 拿到uncle结点
			Entry<K, V> uncleNode = parentNode.getSibling();
			// 拿到祖父节点
			Entry<K, V> grandNode = parentNode.parent;
			// 父结点是红色说明一定有祖父节点
//            if(grandNode == null) {
//                dyeingBlack(parentNode);
//                return;
//            }
			// uncle结点是红色
			if (isRed(uncleNode)) {
				dyeingBlack(parentNode);
				dyeingBlack(uncleNode);
				dyeingRed(grandNode);
				// 把祖父节点当做是新添加的结点
				afterPut(grandNode);
				return;
			}
			// LL RR LR RL
			else {
				if (parentNode.isLeftChildNode()) {
					// LL
					if (node.isLeftChildNode()) {
						dyeingBlack(parentNode);
						dyeingRed(grandNode);

						rotateRight(grandNode);
					}
					// LR
					else {
						dyeingBlack(node);
						dyeingRed(grandNode);

						rotateLeft(parentNode);
						rotateRight(grandNode);
					}
				} else {
					// RL
					if (node.isLeftChildNode()) {
						dyeingBlack(node);
						dyeingRed(grandNode);

						rotateRight(parentNode);
						rotateLeft(grandNode);
					}
					// RR
					else {
						dyeingBlack(parentNode);
						dyeingRed(grandNode);

						rotateLeft(grandNode);
					}
				}

			}
		}
	}

	// 左旋 右边高
	private void rotateLeft(@NotNull Entry<K, V> node) {
		Entry<K, V> childNode = node.right;
		Entry<K, V> childLeftChildNode = childNode.left;
		node.right = childLeftChildNode;
		childNode.left = node;

		afterRotate(childLeftChildNode, childNode, node);
	}

	// 右旋
	private void rotateRight(@NotNull Entry<K, V> node) {
		Entry<K, V> childNode = node.left;
		Entry<K, V> childRightChildNode = childNode.right;
		node.left = childRightChildNode;
		childNode.right = node;

		afterRotate(childRightChildNode, childNode, node);
	}

	private void afterRotate(Entry<K, V> childRightChildNode, Entry<K, V> childNode, Entry<K, V> node) {
		if (childRightChildNode != null) childRightChildNode.parent = node;
		childNode.parent = node.parent;


		if (node.parent == null) root = childNode;
		else if (node.isLeftChildNode()) node.parent.left = childNode;
		else if (node.isRightChildNode()) node.parent.right = childNode;

		node.parent = childNode;

	}

	// 结点染色
	private Entry<K, V> dyeing(Entry<K, V> entry, boolean color) {
		if (entry == null) return null;
		entry.color = color;
		return entry;
	}

	private Entry<K, V> dyeingRed(Entry<K, V> node) {
		return dyeing(node, RED);
	}

	private Entry<K, V> dyeingBlack(Entry<K, V> node) {
		return dyeing(node, BLACK);
	}

	// 判断是什么颜色
	private boolean colorOf(Entry<K, V> node) {
		return node == null ? BLACK : node.color;
	}

	private boolean isBlack(Entry<K, V> node) {
		return colorOf(node) == BLACK;
	}

	private boolean isRed(Entry<K, V> node) {
		return colorOf(node) == RED;
	}


	private void elementNotNullCheck(K key) {
		if (key == null) throw new IllegalArgumentException("element must not be null");
	}

	private Entry<K, V> createEntry(K key, V value, Entry<K, V> left, Entry<K, V> right, Entry<K, V> parent) {
		return new Entry<>(key, value, left, right, parent);
	}

	private static class Entry<K, V> {
		// 默认RED
		boolean color = RED;
		K key;
		V value;
		Entry<K, V> left;
		Entry<K, V> right;
		Entry<K, V> parent;

		public Entry(K key, V value, Entry<K, V> left, Entry<K, V> right, Entry<K, V> parent) {
			this.key = key;
			this.value = value;
			this.left = left;
			this.right = right;
			this.parent = parent;
		}

		public boolean isLeaf() {
			return this.left == null && this.right == null;
		}

		public boolean hasFullChild() {
			return this.left != null && this.right != null;
		}

		public boolean isLeftChildNode() {
			return parent != null && parent.left == this;
		}

		public boolean isRightChildNode() {
			return parent != null && parent.right == this;
		}

		public Entry<K, V> getSibling() {
			if (this.isLeftChildNode()) return this.parent.right;
			else if (this.isRightChildNode()) return this.parent.left;
			// this.parent == null => this是root结点
			return null;
		}
	}
}
